import asyncio
import requests
from mitmproxy.tools.dump import DumpMaster
from mitmproxy.options import Options
from mitmproxy import http
import re
import json
from urllib.parse import urlparse

# 未授权检测API地址
UNAUTH_CHECK_API = "http://localhost:5000/api/unauth-check"

# 需要忽略的域名（正则表达式字符串）
IGNORED_DOMAINS = [
    r".*baidu\.com.*", r".*firefox\.com.*", r".*bdimg\.com.*",
    r".*github\.com.*", r".*githubusercontent\.com.*", r".*googleapis\.com.*",
    r".*google\.com.*", r".*bdstatic\.com.*", r".*mozilla\.org.*",
    r".*mozilla\.com.*", r".*qq\.com,*", r".*openai\.com.*",
    r".*localhost.*", r".*127\.0\.0\.1.*"  # 添加本地地址避免循环
]

# 需要检测的路径模式（可选）
TARGET_PATHS = [
    r"/api/.*", r"/admin/.*", r"/user/.*", r"/data/.*",
    r"/config/.*", r"/system/.*", r"/manage/.*"
]


class UnauthorizedDetector:
    def __init__(self):
        self.processed_requests = set()  # 避免重复处理
        self.detection_count = 0

    def is_ignored_domain(self, url):
        """检查是否是需要忽略的域名"""
        try:
            domain = urlparse(url).netloc
            for pattern in IGNORED_DOMAINS:
                if re.match(pattern, domain, re.IGNORECASE):
                    return True
            return False
        except:
            return True

    def is_target_path(self, url):
        """检查是否是目标路径"""
        try:
            path = urlparse(url).path
            for pattern in TARGET_PATHS:
                if re.match(pattern, path, re.IGNORECASE):
                    return True
            return False
        except:
            return False

    def request_to_raw_http(self, flow: http.HTTPFlow):
        """将mitmproxy请求转换为原始HTTP格式"""
        try:
            # 构建请求行
            request_line = f"{flow.request.method} {flow.request.path} {flow.request.http_version}"

            # 构建请求头
            headers = []
            for key, value in flow.request.headers.items():
                # 跳过一些不需要的头部
                if key.lower() not in ['content-length', 'host']:
                    headers.append(f"{key}: {value}")

            # 添加Host头（必须）
            headers.append(f"Host: {flow.request.host}")

            # 构建请求体
            body = flow.request.get_text() or ""

            # 组合成原始HTTP请求
            raw_request = f"{request_line}\n" + "\n".join(headers) + "\n\n" + body

            return raw_request
        except Exception as e:
            print(f"Error converting request to raw HTTP: {e}")
            return None

    def send_to_detection_api(self, raw_request):
        """发送到未授权检测API"""
        try:
            payload = {"raw_request": raw_request}
            response = requests.post(
                UNAUTH_CHECK_API,
                json=payload,
                timeout=10,
                headers={'Content-Type': 'application/json'}
            )

            if response.status_code == 200:
                result = response.json()
                return result
            else:
                print(f"API returned error: {response.status_code} - {response.text}")
                return None

        except requests.exceptions.RequestException as e:
            print(f"Request to detection API failed: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"JSON decode error: {e}")
            return None

    def request(self, flow: http.HTTPFlow):
        """请求拦截逻辑"""
        # 跳过忽略的域名
        if self.is_ignored_domain(flow.request.url):
            return

        # 只处理特定路径（可选）
        # if not self.is_target_path(flow.request.url):
        #     return

        # 避免重复处理同一请求
        request_id = f"{flow.request.method}_{flow.request.url}"
        if request_id in self.processed_requests:
            return
        self.processed_requests.add(request_id)

        # 检查是否有认证信息
        has_auth = any(key.lower() in ['authorization', 'cookie', 'token']
                       for key in flow.request.headers.keys())

        if not has_auth:
            # 如果没有认证信息，跳过检测
            return

        print(f"检测请求: {flow.request.method} {flow.request.url}")

        # 转换为原始HTTP格式
        raw_request = self.request_to_raw_http(flow)
        if not raw_request:
            return

        # 发送到检测API
        result = self.send_to_detection_api(raw_request)

        if result:
            self.detection_count += 1
            if result.get('is_unauthorized', False):
                print(f"🚨 发现未授权访问漏洞! ({self.detection_count})")
                print(f"   URL: {flow.request.url}")
                print(f"   状态码: {result['original_response']['status_code']}")
                print(f"   响应大小: {result['original_response']['size']} bytes")
                print("=" * 50)
            else:
                print(f"✅ 请求需要认证 ({self.detection_count})")


class ProxyRunner(DumpMaster):
    def __init__(self, options):
        super().__init__(options)
        self.detector = UnauthorizedDetector()
        self.addons.add(self.detector)


async def start_proxy():
    options = Options(
        listen_host="0.0.0.0",
        listen_port=8080,
        ignore_hosts=IGNORED_DOMAINS,
        ssl_insecure=True,  # 允许代理忽略SSL证书验证
        http2=False  # 禁用HTTP/2以简化处理
    )

    proxy = ProxyRunner(options)
    print("=" * 60)
    print("🚀 未授权访问检测代理已启动")
    print(f"📡 监听端口: 8080")
    print(f"🔍 检测API: {UNAUTH_CHECK_API}")
    print(f"🚫 忽略域名: {len(IGNORED_DOMAINS)}个")
    print(f"🎯 目标路径: {TARGET_PATHS}")
    print("=" * 60)
    print("等待流量... (Ctrl+C 停止)")

    try:
        await proxy.run()
    except KeyboardInterrupt:
        print("\n🛑 代理服务器已停止")
        print(f"📊 总共检测了 {proxy.detector.detection_count} 个请求")


def test_connection():
    """测试与检测API的连接"""
    try:
        response = requests.get("http://localhost:5000/health", timeout=5)
        if response.status_code == 200:
            print("✅ 检测API连接正常")
            return True
        else:
            print("❌ 检测API返回错误状态码")
            return False
    except requests.exceptions.RequestException:
        print("❌ 无法连接到检测API，请确保服务已启动")
        return False


if __name__ == "__main__":
    # 检查API连接
    if not test_connection():
        print("请先启动检测API服务: python app.py")
        exit(1)

    # 启动代理
    asyncio.run(start_proxy())